#pragma once

#include <glib.h>
#include <assert.h>

namespace BuildNG
{
  /// I'm a game programmer, I freak out about performance :)
  template <int on_stack_size = 1024>
  class TmpAllocator
  {
  public:
    TmpAllocator()
    {
      allocated = 0;
      allocated_size = 0;
    }

    ~TmpAllocator()
    {
      assert(allocated_size == 0);    // check for memory leaks
    }

    void* alloc(size_t size)
    {
      assert(allocated_size == 0);

      if(size == 0)
        size = 1;

      allocated_size = size;
      if(size <= on_stack_size)
        return on_stack;
      else
      {
        allocated = (gchar*)malloc(size);
        if(!allocated)
          throw std::bad_alloc();

        return allocated;
      }
    }

    void* realloc(void *cur,size_t new_size)
    {
      if(cur == 0)
        return alloc(new_size);

      assert(cur == allocated || cur == on_stack);
      assert(allocated_size != 0);

      if(new_size == 0)
        new_size = 1;

      if(allocated)
      {
        if(new_size <= on_stack_size)
        {
          memcpy(on_stack,allocated,new_size);
          free(allocated);
          allocated_size = new_size;
          return on_stack;
        }
        else
        {
          allocated = (gchar*)::realloc(cur,new_size);
          if(!allocated)
            throw std::bad_alloc();

          allocated_size = new_size;
          return allocated;
        }
      }
      else
      {
        if(new_size <= on_stack_size)
        {
          allocated_size = new_size;
          return on_stack;
        }
        else
        {
          allocated = (gchar*)malloc(new_size);
          if(!allocated)
            throw std::bad_alloc();

          memcpy(allocated,on_stack,std::min(new_size,allocated_size));
          allocated_size = new_size;
          return allocated;
        }
      }
    }

    void free(void *ptr)
    {
      if(!ptr)
        return;

      assert(ptr == on_stack || ptr == allocated);

      if(allocated)
      {
        ::free(allocated);
        allocated = 0;
      }

      allocated_size = 0;
    }

  private:
    gchar on_stack[on_stack_size];
    gchar *allocated;
    size_t allocated_size;
  };

  template <int on_stack_size = 1024>
  class TmpPool
  {
  public:
    TmpPool()
    {
      on_stack_alloc_ptr = 0;
      bucket_list_head = 0;
    }

    ~TmpPool()
    {
      clear();
    }

    void* alloc(int size)
    {
      if(on_stack_alloc_ptr + size <= on_stack_size)
      {
        char *ret = on_stack + on_stack_alloc_ptr;
        on_stack_alloc_ptr += size;
        return ret;
      }
      else
      {
        Bucket *cur_bucket = bucket_list_head;
        if(cur_bucket)
        {
          while(1)
          {
            if(cur_bucket->alloc_ptr + size <= cur_bucket->size)
            {
              char *ret = cur_bucket->bytes + cur_bucket->alloc_ptr;
              cur_bucket->alloc_ptr += size;
              return ret;
            }

            if(!cur_bucket->next)
              break;
            cur_bucket = cur_bucket->next;
          }
        }

        // we have to alloc a new bucket
        int new_bucket_size = size;
        if(cur_bucket)
          new_bucket_size = std::max(size,cur_bucket->size * 2);

        Bucket *new_bucket = (Bucket*)malloc(sizeof(Bucket) + new_bucket_size);
        if(!new_bucket)
          throw std::bad_alloc();

        new_bucket->next = 0;
        new_bucket->size = new_bucket_size;
        new_bucket->alloc_ptr = size;

        if(cur_bucket)
          cur_bucket->next = new_bucket;
        else
          bucket_list_head = new_bucket;

        return new_bucket->bytes;
      }
    }

    void clear()
    {
      Bucket *cur = bucket_list_head;
      while(cur)
      {
        Bucket *next = cur->next;
        free(cur);
        cur = next;
      }

      on_stack_alloc_ptr = 0;
      bucket_list_head = 0;
    }

  private:
    char          on_stack[on_stack_size];
    int           on_stack_alloc_ptr;

    struct Bucket
    {
      Bucket*     next;
      
      int         size;
      int         alloc_ptr;
      char        bytes[1];
    };
    Bucket*       bucket_list_head;
  };

  template <class T,int on_stack_size = 1024>
  class TmpArray
  {
    template <class FriendT,int friend_on_stack_size> friend class TmpArray;

  public:
    TmpArray();
    TmpArray(const TmpArray<T,on_stack_size> &src);
    template <int other_size>
    TmpArray(const TmpArray<T,other_size> &src);
    ~TmpArray();

    TmpArray<T,on_stack_size>& operator = (const TmpArray<T,on_stack_size> &src);
    template <int other_size>
    TmpArray<T,on_stack_size>& operator = (const TmpArray<T,other_size> &src);


    void      append(const T &t);
    void      remove(int index);
    void      clear();

    int       get_size() const;

    T&        operator [] (guint i);
    const T&  operator [] (guint i) const;

    T&        back();
    const T&  back() const;

    void      replace(int begin,int end,const T &new_val);

    T*        get_begin() { return data; }
    T*        get_end() { return data + size; }

  private:
    size_t    reserved;
    size_t    size;
    T*        data;

    char      pre_allocated[on_stack_size * sizeof(T)];
    T*        allocated;

    template <int other_size>
    void      copy(const TmpArray<T,other_size> &src);
    void      clean_up();
  };

  /////////
  // TmpArray implementation
  template <class T,int on_stack_size>
  TmpArray<T,on_stack_size>::TmpArray()
  {
    reserved = on_stack_size;
    size = 0;
    data = (T*)pre_allocated;
    allocated = 0;
  }

  template <class T,int on_stack_size>
  TmpArray<T,on_stack_size>::TmpArray(const TmpArray<T,on_stack_size> &src)
  {
    size = reserved = 0;
    data = 0;
    allocated = 0;

    copy(src);
  }

  template <class T,int on_stack_size>
  template <int other_size>
  TmpArray<T,on_stack_size>::TmpArray(const TmpArray<T,other_size> &src)
  {
    size = reserved = 0;
    data = 0;
    allocated = 0;

    copy(src);
  }

  template <class T,int on_stack_size>
  TmpArray<T,on_stack_size>::~TmpArray()
  {
    clean_up();
  }

  template <class T,int on_stack_size>
  TmpArray<T,on_stack_size>& TmpArray<T,on_stack_size>::operator = (const TmpArray<T,on_stack_size> &src)
  {
    copy(src);
    return *this;
  }

  template <class T,int on_stack_size>
  template <int other_size>
  TmpArray<T,on_stack_size>& TmpArray<T,on_stack_size>::operator = (const TmpArray<T,other_size> &src)
  {
    copy(src);
    return *this;
  }

  template <class T,int on_stack_size>
  void TmpArray<T,on_stack_size>::append(const T &t)
  {
    if(size >= reserved)
    {
      size_t new_reserved = reserved * 2;
      T *new_allocated = (T*)malloc(new_reserved * sizeof(T));
      if(!new_allocated)
        throw std::bad_alloc();

      int num_copied = 0;
      try
      {
        for(size_t i = 0;i < size;i++)
        {
          new (new_allocated + i) T(data[i]);
          num_copied++;
        }
      }
      catch(...)
      {
        for(int i = 0;i < num_copied;i++)
        {
          for(size_t i = 0;i < size;i++)
          {
            T *cur = new_allocated + i;
            cur->~T();
          }
        }
        free(new_allocated);
        throw;
      }

      for(size_t i = 0;i < size;i++)
      {
        T *cur = data + i;
        cur->~T();
      }

      free(allocated);    // if allocated wasn't used yet, this is still 0.
      allocated = new_allocated;
      reserved = new_reserved;
      data = allocated;
    }

    new(data + size) T(t);
    size++;
  }

  template <class T,int on_stack_size>
  void TmpArray<T,on_stack_size>::remove(int index)
  {
    assert((size_t)index < size);

    if(size - 1 < reserved / 2 && allocated)
    {
      // we can shrink

      size_t new_reserved = reserved / 2;
      T *new_data;
      T *new_allocated;
      if(new_reserved > on_stack_size)
      {
        new_allocated = (T*)malloc(new_reserved * sizeof(T));
        if(!new_allocated)
          throw std::bad_alloc();
        new_data = new_allocated;
        new_reserved = on_stack_size;
      }
      else
      {
        new_data = (T*)pre_allocated;
        new_reserved = on_stack_size;
        new_allocated = 0;
      }

      size_t num_copied = 0;

      try
      {
        for(int i = 0;i < index;i++)
        {
          new (&new_data[i]) T(data[i]);
          num_copied++;
        }

        for(size_t i = index + 1;i < size;i++)
        {
          new (&new_data[i - 1]) T(data[i]);
          num_copied++;
        }
      }
      catch(...)
      {
        for(size_t i = 0;i < num_copied;i++)
        {
          new_data[i].~T();
          free(new_allocated);
        }
        throw;
      }

      for(size_t i = 0;i < size;i++)
      {
        data[i].~T();
      }

      data = new_data;
      reserved = new_reserved;
      allocated = new_allocated;

      size--;
    }
    else
    {
      data[index].~T();
      int free_slot = index;
      try
      {
        size--;

        for(size_t i = index;i < size;i++)
        {
          new (data + i) T(data[i + 1]);
          data[i + 1].~T();

          free_slot++;
        }
      }
      catch(...)
      {
        for(size_t i = free_slot;i < size;i++)
        {
          data[i].~T();
        }
        size = free_slot;
        throw;
      }
    }
  }

  template <class T,int on_stack_size>
  int TmpArray<T,on_stack_size>::get_size() const
  {
    return (int)size;
  }

  template <class T,int on_stack_size>
  T& TmpArray<T,on_stack_size>::operator [] (guint i)
  {
    assert(i < size);
    return data[i];
  }

  template <class T,int on_stack_size>
  const T& TmpArray<T,on_stack_size>::operator [] (guint i) const
  {
    assert(i < size);
    return data[i];
  }

  template <class T,int on_stack_size>
  T& TmpArray<T,on_stack_size>::back()
  {
    assert(size != 0);
    return data[size - 1];
  }

  template <class T,int on_stack_size>
  const T& TmpArray<T,on_stack_size>::back() const
  {
    assert(size != 0)
    return data[size - 1];
  }

  template <class T,int on_stack_size>
  template <int other_size>
  void TmpArray<T,on_stack_size>::copy(const TmpArray<T,other_size> &src)
  {
    clean_up();

    reserved = std::max<size_t>(on_stack_size,src.size);
    size = src.size;
    if(reserved > on_stack_size)
    {
      allocated = (T*)malloc(reserved * sizeof(T));
      if(!allocated)
        throw std::bad_alloc();
      data = allocated;
    }
    else
    {
      data = (T*)pre_allocated;
    }

    for(int i = 0;i < (int)size;i++)
    {
      try
      {
        new (data + i) T(src[i]);
      }
      catch(...)
      {
        size = i;
        throw;
      }
    }
  }

  template <class T,int on_stack_size>
  void TmpArray<T,on_stack_size>::clean_up()
  {
    for(size_t i = 0;i < size;i++)
    {
      T *cur = data + i;
      cur->~T();
    }

    free(allocated);
    allocated = 0;

    size = reserved = 0;
    data = 0;
  }

  template <class T,int on_stack_size>
  void TmpArray<T,on_stack_size>::clear()
  {
    // we can keep the memory blob reserved, but we do have to
    // call the dtors of the items though.

    for(size_t i = 0;i < size;i++)
    {
      T *cur = data + i;
      cur->~T();
    }

    size = 0;
  }

  template <class T,int on_stack_size>
  void TmpArray<T,on_stack_size>::replace(int begin,int end,const T &new_val)
  {
    assert((size_t)begin < size);
    assert((size_t)end <= size);
    assert(begin <= end);

    if(begin != end)
    {
      T *cur = data + begin;
      cur->~T();
      new(cur) T(new_val);

      size_t move = end - begin - 1;

      if(move != 0)
      {
        for(size_t i = begin + 1;i < size - move;i++)
        {
          T *cur = data + i;
          cur->~T();

          new(cur) T(*(data + i + move));
        }

        for(size_t i = size - move;i < size;i++)
        {
          T *cur = data + i;
          cur->~T();
        }

        size -= move;
      }
    }
    else
    {
      assert(!"not implemented yet");
    }
  }
}

